from DrissionPage import ChromiumPage, ChromiumOptions
from DrissionPage.configs import options_manage
import tkinter.filedialog as tf
from pathlib import Path
import tqdm

line_count = input('每张图下载几张：')
try:
    line_count = int(line_count)
    if 0 < line_count <= 30:
        pass
    else:
        input('请输入1-30的数字')
        exit()
except:
    input('请输入数字')
    exit()
path = Path(tf.askdirectory(title='选择文件夹'))
imgs = list(path.rglob('*.jpg')) + list(path.rglob('*.png'))
co = ChromiumOptions(read_file=False)
co.set_argument('--incognito')
co.set_argument('--headless')
co.set_paths(browser_path=tf.askopenfilename(title='选择浏览器路径'))
page = ChromiumPage(addr_driver_opts=co)

for file_path in tqdm.tqdm(imgs):
    page.get('https://graph.baidu.com/pcpage/index?tpl_from=pc')

    file_input = page.ele('xpath://input[@name="file"]')

    file_input.input(str(file_path))

    urls = page.eles('xpath://div[@class="general-waterfall"]//img/@src')[:line_count]

    for index, item in enumerate(urls):
        save_path = Path(f'result/{file_path.stem}/{index}.jpg')
        save_path.parent.mkdir(parents=True, exist_ok=True)
        with open(save_path, 'wb') as f:
            f.write(item.get_src())

input('完成...')
